import base64
from zhipuai import ZhipuAI
from API_KEY import *

client = ZhipuAI(api_key=api_key_zhipuai)  # 填写您自己的APIKey


def image_compre(img_path):
    with open(img_path, 'rb') as img_file:
        img_base = base64.b64encode(img_file.read()).decode('utf-8')

    response = client.chat.completions.create(
        model="glm-4v-plus-0111",  # 填写需要调用的模型名称
        messages=[
            {
                "role": "user",
                "content": [
                    {
                        "type": "image_url",
                        "image_url": {
                            "url": img_base
                        }
                    },
                    {
                        "type": "text",
                        "text": "解释图中的内容"
                    }
                ]
            }
        ]
    )
    return response.choices[0].message.content


def draw(prompt):
    response = client.images.generations(
        model="cogview-4-250304",  # 填写需要调用的模型编码
        prompt=prompt,
    )
    return response.data[0].url


if __name__ == '__main__':
    # 图像理解
    # res = image_compre("xiongmao.jpeg")
    # print(res)
    # 图像生成
    res = draw("一只大熊猫，坐在山林一棵高高的树上")
    print(res)
